#!/usr/bin/env python3

import asyncio
import gc
import logging
import re
import time
from concurrent.futures import ThreadPoolExecutor
from typing import Any, Dict, List, Optional

# 尝试导入不同的依赖
try:
    from modelscope.pipelines import pipeline
    from modelscope.utils.config import Config
    MODELSCOPE_AVAILABLE = True
    FUNASR_AVAILABLE = False
except ImportError:
    try:
        from funasr import AutoModel
        FUNASR_AVAILABLE = True
        MODELSCOPE_AVAILABLE = False
        logging.warning("ModelScope未安装，但FunASR可用。注意：FunASR不支持NLP模型")
    except ImportError:
        logging.warning("ModelScope和FunASR都未安装，请先安装其中之一")
        MODELSCOPE_AVAILABLE = False
        FUNASR_AVAILABLE = False

from ...config.funasr_config import ModelConfig
from ..schemas import ActionClassification, SummaryResult

logger = logging.getLogger(__name__)

class Palm2Manager:
    """
    PALM2.0模型管理器（性能优化版本）

    负责管理PALM2.0模型的加载、推理和资源管理
    用于生成会议纪要和文本摘要
    包含性能优化和内存管理功能
    """

    def __init__(self):
        self.model = None
        self.model_loaded = False
        self.model_name = ModelConfig.SUMMARY_MODEL  # 使用large版本
        self.model_revision = ModelConfig.SUMMARY_MODEL_REVISION
        self.load_time = None
        self.last_inference_time = None

        # 线程池优化
        self.executor = ThreadPoolExecutor(
            max_workers=3,
            thread_name_prefix="palm2"
        )

        # 性能统计
        self.stats = {
            'total_inferences': 0,
            'total_inference_time': 0,
            'average_inference_time': 0,
            'load_attempts': 0,
            'successful_loads': 0,
            'errors': 0,
            'cache_hits': 0,
            'cache_misses': 0
        }

        # 生成配置（优化后）
        self.generation_config = {
            'max_length': 512,
            'min_length': 50,
            'temperature': 0.7,
            'top_p': 0.9,
            'repetition_penalty': 1.1,
            'do_sample': True,
            'num_return_sequences': 1,
            'early_stopping': True
        }

        # 模型配置（性能优化）
        self.config = {
            'model': self.model_name,
            'model_revision': self.model_revision,
            'device': 'cpu',  # 可以根据配置调整为'gpu'
        }

        # 摘要模板和配置
        self.summary_templates = {
            'meeting_minutes': """
请根据以下会议内容生成简洁的会议纪要：

原文：{text}

{classification_info}

请按以下格式生成纪要：
1. 会议要点
2. 重要决策
3. 待办事项
4. 后续跟进

纪要：""",
            'simple_summary': """
请为以下文本生成简洁摘要：

{text}

摘要：""",
            'detailed_summary': """
请为以下文本生成详细摘要，包含关键信息和要点：

{text}

{classification_info}

详细摘要："""
        }

        # 文本清理配置
        self.text_cleaning = {
            'remove_extra_spaces': True,
            'remove_special_chars': True,
            'normalize_punctuation': True,
            'max_summary_length': 800,
            'min_summary_length': 20
        }

    async def load_model(self) -> bool:
        """
        异步加载模型（性能优化版本）

        Returns:
            bool: 加载是否成功
        """
        if self.model_loaded and self.model is not None:
            logger.info("PALM2.0模型已加载")
            return True

        # 优先使用ModelScope（因为这是NLP模型）
        if MODELSCOPE_AVAILABLE:
            return await self._load_with_modelscope()
        else:
            logger.error("ModelScope未安装，无法加载NLP模型")
            logger.info("提示：请安装ModelScope：pip install modelscope")
            return False

    async def _load_with_modelscope(self) -> bool:
        """使用ModelScope加载模型"""

        self.stats['load_attempts'] += 1
        start_time = time.time()

        try:
            logger.info(f"开始加载PALM2.0模型: {self.model_name}")

            # 使用新的模型加载器
            from .modelscope_loader import get_modelscope_loader
            loader = get_modelscope_loader()

            # 在线程池中加载模型以避免阻塞
            loop = asyncio.get_event_loop()
            self.model = await loop.run_in_executor(
                self.executor,
                self._load_model_sync_modelscope,
                loader
            )

            if self.model is not None:
                self.model_loaded = True
                self.load_time = time.time() - start_time
                self.stats['successful_loads'] += 1
                logger.info(f"PALM2.0模型加载成功，耗时: {self.load_time:.2f}秒")
                return True
            else:
                logger.error("PALM2.0模型加载失败")
                return False

        except Exception as e:
            self.stats['errors'] += 1
            logger.error(f"PALM2.0模型加载异常: {e}")
            return False

    def _load_model_sync_modelscope(self, loader):
        """使用ModelScope同步加载模型"""
        try:
            # 使用新的加载器从本地加载模型
            model = loader.load_model_with_fallback(
                model_type="palm2",
                task='text-generation',
                device=self.config['device']
            )
            return model
        except Exception as e:
            logger.error(f"ModelScope同步加载模型失败: {e}")
            return None

    async def generate_summary(
        self,
        text: str,
        classification_result: Optional[ActionClassification] = None,
        summary_type: str = "meeting_minutes",
        custom_prompt: Optional[str] = None
    ) -> Optional[SummaryResult]:
        """
        生成文本摘要（性能优化版本）

        Args:
            text: 输入文本
            classification_result: 分类结果
            summary_type: 摘要类型
            custom_prompt: 自定义提示词

        Returns:
            Optional[SummaryResult]: 摘要结果
        """
        if not self.model_loaded or self.model is None:
            logger.warning("模型未加载，尝试加载模型")
            if not await self.load_model():
                logger.error("模型加载失败，无法生成摘要")
                return None

        start_time = time.time()

        try:
            # 构建提示词
            prompt = self._build_prompt(text, classification_result, summary_type, custom_prompt)

            # 检查文本长度和预处理
            processed_text = self._preprocess_text(text)
            if not processed_text:
                logger.warning("预处理后文本为空")
                return self._generate_fallback_summary(text)

            # 执行生成
            summary_text = await self._generate_text(prompt)

            if not summary_text:
                logger.warning("生成的摘要为空，使用备用方案")
                return self._generate_fallback_summary(text)

            # 后处理和提取关键信息
            processed_summary = self._postprocess_summary(summary_text)
            key_points = self._extract_key_points(processed_summary, classification_result)

            # 构建结果
            result = SummaryResult(
                summary=processed_summary,
                key_points=key_points,
                word_count=len(processed_summary),
                generation_time=time.time() - start_time,
                model_name=self.model_name
            )

            # 更新统计信息
            inference_time = time.time() - start_time
            self.last_inference_time = inference_time
            self.stats['total_inferences'] += 1
            self.stats['total_inference_time'] += inference_time
            self.stats['average_inference_time'] = (
                self.stats['total_inference_time'] / self.stats['total_inferences']
            )

            logger.info(f"摘要生成完成，耗时: {inference_time:.3f}秒，字数: {len(processed_summary)}")
            return result

        except Exception as e:
            self.stats['errors'] += 1
            logger.error(f"摘要生成失败: {e}")
            return self._generate_fallback_summary(text)

    def _build_prompt(
        self,
        text: str,
        classification_result: Optional[ActionClassification] = None,
        summary_type: str = "meeting_minutes",
        custom_prompt: Optional[str] = None
    ) -> str:
        """
        构建优化的提示词

        Args:
            text: 输入文本
            classification_result: 分类结果
            summary_type: 摘要类型
            custom_prompt: 自定义提示词

        Returns:
            str: 构建的提示词
        """
        if custom_prompt:
            return custom_prompt.format(text=text)

        # 构建分类信息
        classification_info = ""
        if classification_result:
            classification_info = f"\n文本类型：{classification_result.label}（置信度：{classification_result.confidence:.2f}）"

        # 选择模板
        template = self.summary_templates.get(summary_type, self.summary_templates['meeting_minutes'])

        return template.format(
            text=text,
            classification_info=classification_info
        ).strip()

    def _preprocess_text(self, text: str) -> str:
        """
        文本预处理

        Args:
            text: 输入文本

        Returns:
            str: 预处理后的文本
        """
        if not text:
            return ""

        processed = text.strip()

        if self.text_cleaning['remove_extra_spaces']:
            processed = re.sub(r'\s+', ' ', processed)

        if self.text_cleaning['remove_special_chars']:
            # 保留常用标点，移除特殊字符
            processed = re.sub(r'[^\w\s\u4e00-\u9fff。，！？；：""''（）【】[\].,!?;:()\-]', '', processed)

        if self.text_cleaning['normalize_punctuation']:
            # 标准化标点符号
            processed = processed.replace('，', ',').replace('。', '.').replace('！', '!').replace('？', '?')

        return processed

    async def _generate_text(self, prompt: str) -> Optional[str]:
        """
        异步生成文本

        Args:
            prompt: 提示词

        Returns:
            Optional[str]: 生成的文本
        """
        try:
            # 在线程池中执行生成
            loop = asyncio.get_event_loop()
            result = await loop.run_in_executor(
                self.executor,
                self._generate_text_sync,
                prompt
            )

            if result and isinstance(result, list) and len(result) > 0:
                # ModelScope pipeline返回格式：[{'generated_text': '...'}]
                generated_text = result[0].get('generated_text', '')

                # 移除提示词部分，只保留生成的内容
                if prompt in generated_text:
                    generated_text = generated_text.replace(prompt, '').strip()

                return generated_text
            else:
                logger.warning(f"生成结果格式异常: {result}")
                return None

        except Exception as e:
            logger.error(f"文本生成失败: {e}")
            return None

    def _generate_text_sync(self, prompt: str):
        """同步生成文本的内部方法"""
        try:
            # 限制提示词长度
            max_prompt_length = 800
            if len(prompt) > max_prompt_length:
                prompt = prompt[:max_prompt_length] + "..."

            # 使用ModelScope pipeline的调用方式
            result = self.model(
                prompt,
                **self.generation_config
            )
            return result
        except Exception as e:
            logger.error(f"同步生成文本失败: {e}")
            return None

    def _postprocess_summary(self, summary_text: str) -> str:
        """
        后处理摘要文本

        Args:
            summary_text: 原始摘要文本

        Returns:
            str: 处理后的摘要文本
        """
        if not summary_text:
            return ""

        # 基本清理
        processed = summary_text.strip()

        # 移除重复内容
        processed = self._remove_duplicates(processed)

        # 长度控制
        max_length = self.text_cleaning['max_summary_length']
        if len(processed) > max_length:
            # 在句号处截断
            sentences = processed.split('。')
            truncated = ""
            for sentence in sentences:
                if len(truncated + sentence + '。') <= max_length:
                    truncated += sentence + '。'
                else:
                    break
            processed = truncated or processed[:max_length]

        # 最小长度检查
        min_length = self.text_cleaning['min_summary_length']
        if len(processed) < min_length:
            processed += f" （摘要长度：{len(processed)}字符）"

        return processed

    def _remove_duplicates(self, text: str) -> str:
        """
        移除重复句子

        Args:
            text: 输入文本

        Returns:
            str: 去重后的文本
        """
        sentences = text.split('。')
        unique_sentences = []
        seen = set()

        for sentence in sentences:
            sentence = sentence.strip()
            if sentence and sentence not in seen:
                unique_sentences.append(sentence)
                seen.add(sentence)

        return '。'.join(unique_sentences)

    def _extract_key_points(
        self,
        summary: str,
        classification_result: Optional[ActionClassification] = None
    ) -> List[str]:
        """
        从摘要中提取关键点

        Args:
            summary: 摘要文本
            classification_result: 分类结果

        Returns:
            List[str]: 关键点列表
        """
        key_points = []

        # 基于数字标记提取
        numbered_items = re.findall(r'[0-9]+[\.、]\s*([^0-9\n]+)', summary)
        key_points.extend([item.strip() for item in numbered_items if item.strip()])

        # 基于关键词提取
        keywords = ['决策', '任务', '问题', '重要', '关键', '核心', '主要']
        sentences = summary.split('。')

        for sentence in sentences:
            sentence = sentence.strip()
            if any(keyword in sentence for keyword in keywords) and len(sentence) > 10:
                key_points.append(sentence)

        # 基于分类结果添加关键点
        if classification_result:
            if classification_result.label == "决策" and classification_result.confidence > 0.7:
                decision_point = f"重要决策内容（置信度：{classification_result.confidence:.2f}）"
                key_points.insert(0, decision_point)
            elif classification_result.label == "任务" and classification_result.confidence > 0.7:
                task_point = f"任务安排要点（置信度：{classification_result.confidence:.2f}）"
                key_points.insert(0, task_point)

        # 去重和限制数量
        unique_points = []
        seen = set()
        for point in key_points:
            if point not in seen and len(point) > 5:
                unique_points.append(point)
                seen.add(point)
                if len(unique_points) >= 5:  # 最多5个关键点
                    break

        return unique_points

    def _generate_fallback_summary(self, text: str) -> SummaryResult:
        """
        生成备用摘要（当模型推理失败时使用）

        Args:
            text: 输入文本

        Returns:
            SummaryResult: 备用摘要结果
        """
        logger.info("生成备用摘要")

        # 简单的摘要生成逻辑
        sentences = text.replace('。', '.').replace('！', '!').replace('？', '?').split('.')
        important_sentences = [s.strip() for s in sentences if len(s.strip()) > 15]

        if len(important_sentences) > 3:
            summary = "。".join(important_sentences[:3]) + "。"
        elif important_sentences:
            summary = "。".join(important_sentences) + "。"
        else:
            summary = f"文本摘要：{text[:200]}{'...' if len(text) > 200 else ''}"

        # 提取简单关键点
        key_points = []
        for sentence in important_sentences[:5]:
            if len(sentence.strip()) > 20:
                key_points.append(sentence.strip())

        if not key_points:
            key_points = ["无法提取关键点，请查看完整摘要"]

        return SummaryResult(
            summary=summary,
            key_points=key_points,
            word_count=len(summary),
            generation_time=0.0,
            model_name="fallback"
        )

    async def get_model_status(self) -> Dict[str, Any]:
        """
        获取模型状态（增强版本）

        Returns:
            Dict[str, Any]: 模型状态信息
        """
        try:
            return {
                'model_name': self.model_name,
                'model_revision': self.model_revision,
                'loaded': self.model_loaded,
                'load_time': self.load_time,
                'last_inference_time': self.last_inference_time,
                'statistics': self.stats.copy(),
                'config': self.config.copy(),
                'generation_config': self.generation_config.copy(),
                'available_templates': list(self.summary_templates.keys()),
                'modelscope_available': MODELSCOPE_AVAILABLE
            }
        except Exception as e:
            logger.error(f"获取模型状态失败: {e}")
            return {
                'loaded': False,
                'error': str(e)
            }

    async def warm_up(self, sample_texts: Optional[List[str]] = None) -> bool:
        """
        模型预热

        Args:
            sample_texts: 样本文本列表

        Returns:
            bool: 预热是否成功
        """
        if not self.model_loaded:
            logger.warning("模型未加载，无法预热")
            return False

        try:
            # 使用默认样本文本进行预热
            if not sample_texts:
                sample_texts = [
                    "今天的会议讨论了项目进展和下一步计划。",
                    "团队决定采用新的技术方案来提高效率。",
                    "需要在下周五之前完成代码review。"
                ]

            logger.info("开始模型预热...")
            start_time = time.time()

            # 并发执行预热推理
            tasks = [
                self.generate_summary(text, summary_type="simple_summary")
                for text in sample_texts
            ]

            results = await asyncio.gather(*tasks, return_exceptions=True)

            successful_warm_ups = sum(
                1 for r in results
                if isinstance(r, SummaryResult)
            )

            warm_up_time = time.time() - start_time
            logger.info(f"模型预热完成，成功{successful_warm_ups}/{len(sample_texts)}次，耗时: {warm_up_time:.2f}秒")

            return successful_warm_ups > 0

        except Exception as e:
            logger.error(f"模型预热失败: {e}")
            return False

    async def cleanup(self):
        """
        清理资源（增强版本）
        """
        logger.info("开始清理PALM2.0资源...")

        try:
            # 清理模型
            if hasattr(self, 'model') and self.model is not None:
                del self.model
                self.model = None

            self.model_loaded = False

            # 关闭线程池
            if hasattr(self, 'executor') and self.executor:
                self.executor.shutdown(wait=True)

            # 执行垃圾回收
            collected = gc.collect()
            logger.info(f"PALM2.0资源清理完成，回收了{collected}个对象")

        except Exception as e:
            logger.error(f"清理PALM2.0资源时出错: {e}")

    def reset_stats(self):
        """重置统计信息"""
        self.stats = {
            'total_inferences': 0,
            'total_inference_time': 0,
            'average_inference_time': 0,
            'load_attempts': 0,
            'successful_loads': 0,
            'errors': 0,
            'cache_hits': 0,
            'cache_misses': 0
        }
        logger.info("PALM2.0统计信息已重置")

    def update_config(self, new_config: Dict[str, Any]):
        """
        更新配置

        Args:
            new_config: 新配置
        """
        old_config = self.config.copy()
        self.config.update(new_config)
        logger.info(f"PALM2.0配置已更新: {new_config}")

        # 如果设备配置改变，需要重新加载模型
        if old_config.get('device') != self.config.get('device'):
            logger.info("设备配置改变，需要重新加载模型")
            self.model_loaded = False
            self.model = None

    def update_generation_config(self, new_config: Dict[str, Any]):
        """
        更新生成配置

        Args:
            new_config: 新生成配置
        """
        old_config = self.generation_config.copy()
        self.generation_config.update(new_config)
        logger.info(f"PALM2.0生成配置已更新: {new_config}")

    def add_custom_template(self, template_name: str, template_content: str):
        """
        添加自定义摘要模板

        Args:
            template_name: 模板名称
            template_content: 模板内容
        """
        self.summary_templates[template_name] = template_content
        logger.info(f"已添加自定义模板: {template_name}")

    def get_performance_metrics(self) -> Dict[str, Any]:
        """获取性能指标"""
        return {
            'inference_count': self.stats['total_inferences'],
            'average_inference_time': self.stats['average_inference_time'],
            'error_rate': self.stats['errors'] / max(1, self.stats['total_inferences']) * 100,
            'load_success_rate': self.stats['successful_loads'] / max(1, self.stats['load_attempts']) * 100,
            'last_inference_time': self.last_inference_time,
            'model_loaded': self.model_loaded
        }
